package com.wechat.gateway.filter;

import cn.hutool.core.util.StrUtil;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.wechat.gateway.config.WhitelistConfig;
import com.wechat.gateway.util.CryptoUtil;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.reactivestreams.Publisher;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.factory.AbstractGatewayFilterFactory;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpMethod;
import org.springframework.http.MediaType;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.Map;

/**
 * 加密解密过滤器
 * 对请求体进行解密，对响应体进行加密
 */
@Slf4j
@Component
public class EncodeFilter extends AbstractGatewayFilterFactory<EncodeFilter.Config> {

    @Autowired
    private CryptoUtil cryptoUtil;

    @Autowired
    private WhitelistConfig whitelistConfig;

    @Autowired
    private ObjectMapper objectMapper;

    public EncodeFilter() {
        super(Config.class);
    }

    @Override
    public GatewayFilter apply(Config config) {
        return (exchange, chain) -> {
            ServerHttpRequest request = exchange.getRequest();
            ServerHttpResponse response = exchange.getResponse();
            String path = request.getURI().getPath();

            log.debug("加密解密过滤器处理请求: {}", path);

            // 检查是否在加密白名单中
            if (whitelistConfig.isEncryptWhitelist(path)) {
                log.debug("路径 {} 在加密白名单中，跳过加密解密", path);
                return chain.filter(exchange);
            }

            // 只处理POST和PUT请求的请求体
            HttpMethod method = request.getMethod();
            if (HttpMethod.POST.equals(method) || HttpMethod.PUT.equals(method)) {
                // 处理请求体解密
                return handleRequestDecryption(exchange, chain);
            } else {
                // 其他请求只处理响应体加密
                return handleResponseEncryption(exchange, chain);
            }
        };
    }

    /**
     * 处理请求体解密
     */
    private Mono<Void> handleRequestDecryption(org.springframework.web.server.ServerWebExchange exchange,
                                               org.springframework.cloud.gateway.filter.GatewayFilterChain chain) {
        ServerHttpRequest request = exchange.getRequest();
        // 检查Content-Type
        MediaType contentType = request.getHeaders().getContentType();
        if (contentType == null || !MediaType.APPLICATION_JSON.isCompatibleWith(contentType)) {
            log.debug("请求Content-Type不是JSON，跳过解密: {}", contentType);
            return handleResponseEncryption(exchange, chain);
        }
        return DataBufferUtils.join(request.getBody())
                .flatMap(dataBuffer -> {
                    byte[] bytes = new byte[dataBuffer.readableByteCount()];
                    dataBuffer.read(bytes);
                    DataBufferUtils.release(dataBuffer);
                    String body = new String(bytes, StandardCharsets.UTF_8);
                    log.debug("原始请求体: {}", body);

                    // 解密请求体
                    String decryptedBody = decryptRequestBody(body);
                    log.debug("解密后请求体: {}", decryptedBody);

                    // 创建新的请求
                    ServerHttpRequest modifiedRequest = new ServerHttpRequestDecorator(request) {
                        @Override
                        public Flux<DataBuffer> getBody() {
                            DataBuffer buffer = exchange.getResponse().bufferFactory()
                                    .wrap(decryptedBody.getBytes(StandardCharsets.UTF_8));
                            return Flux.just(buffer);
                        }

                        @Override
                        public HttpHeaders getHeaders() {
                            HttpHeaders headers = new HttpHeaders();
                            headers.putAll(super.getHeaders());
                            headers.setContentLength(decryptedBody.getBytes(StandardCharsets.UTF_8).length);
                            return headers;
                        }
                    };

                    return handleResponseEncryption(
                            exchange.mutate().request(modifiedRequest).build(),
                            chain
                    );
                })
                .onErrorResume(throwable -> {
                    log.error("请求体解密失败: {}", throwable.getMessage(), throwable);
                    return handleResponseEncryption(exchange, chain);
                });
    }

    /**
     * 处理响应体加密
     */
    private Mono<Void> handleResponseEncryption(org.springframework.web.server.ServerWebExchange exchange,
                                                org.springframework.cloud.gateway.filter.GatewayFilterChain chain) {
        ServerHttpResponse response = exchange.getResponse();

        ServerHttpResponseDecorator decoratedResponse = new ServerHttpResponseDecorator(response) {
            @Override
            public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
                if (body instanceof Flux) {
                    Flux<? extends DataBuffer> fluxBody = (Flux<? extends DataBuffer>) body;

                    return super.writeWith(fluxBody.map(dataBuffer -> {
                        // 检查响应Content-Type
                        MediaType contentType = response.getHeaders().getContentType();
                        if (contentType == null || !MediaType.APPLICATION_JSON.isCompatibleWith(contentType)) {
                            log.debug("响应Content-Type不是JSON，跳过加密: {}", contentType);
                            return dataBuffer;
                        }

                        byte[] content = new byte[dataBuffer.readableByteCount()];
                        dataBuffer.read(content);

                        String responseBody = new String(content, StandardCharsets.UTF_8);
                        log.debug("原始响应体: {}", responseBody);

                        // 加密响应体
                        String encryptedBody = encryptResponseBody(responseBody);
                        log.debug("加密后响应体: {}", encryptedBody);

                        // 释放原始buffer
                        DataBufferUtils.release(dataBuffer);

                        // 创建新的buffer
                        byte[] encryptedBytes = encryptedBody.getBytes(StandardCharsets.UTF_8);
                        return response.bufferFactory().wrap(encryptedBytes);
                    }));
                }
                return super.writeWith(body);
            }
        };

        return chain.filter(exchange.mutate().response(decoratedResponse).build());
    }

    /**
     * 解密请求体
     */
    private String decryptRequestBody(String encryptedBody) {
        if (StrUtil.isEmpty(encryptedBody)) {
            return encryptedBody;
        }

        try {
            // 尝试解析JSON以检查是否有加密字段
            Map<String, Object> bodyMap = objectMapper.readValue(encryptedBody, Map.class);

            // 检查是否有加密标识
            if (bodyMap.containsKey("encrypted") && Boolean.TRUE.equals(bodyMap.get("encrypted"))) {
                Object dataObj = bodyMap.get("data");
                if (dataObj instanceof String encryptedData) {
                    String decryptedData = cryptoUtil.decrypt(encryptedData);
                    log.debug("解密数据: {} -> {}", encryptedData, decryptedData);
                    return decryptedData;
                }
            }

            // 如果没有加密标识，直接返回原始内容
            return encryptedBody;

        } catch (JsonProcessingException e) {
            log.warn("解析请求体JSON失败，尝试直接解密: {}", e.getMessage());
            try {
                return cryptoUtil.decrypt(encryptedBody);
            } catch (Exception ex) {
                log.warn("直接解密失败，返回原始内容: {}", ex.getMessage());
                return encryptedBody;
            }
        } catch (Exception e) {
            log.error("请求体解密失败: {}", e.getMessage(), e);
            return encryptedBody;
        }
    }

    /**
     * 加密响应体
     */
    private String encryptResponseBody(String responseBody) {
        if (StrUtil.isEmpty(responseBody)) {
            return responseBody;
        }

        try {
            // 加密整个响应体
            String encryptedData = cryptoUtil.encrypt(responseBody);

            // 构造加密响应格式
            Map<String, Object> encryptedResponse = new HashMap<>();
            encryptedResponse.put("encrypted", true);
            encryptedResponse.put("data", encryptedData);
            encryptedResponse.put("timestamp", System.currentTimeMillis());

            return objectMapper.writeValueAsString(encryptedResponse);

        } catch (Exception e) {
            log.error("响应体加密失败: {}", e.getMessage(), e);
            return responseBody;
        }
    }

    /**
     * 配置类
     */
    @Data
    public static class Config {
        /**
         * 是否启用加密解密
         */
        private boolean enabled = true;

        /**
         * 是否启用请求体解密
         */
        private boolean decryptRequest = true;

        /**
         * 是否启用响应体加密
         */
        private boolean encryptResponse = true;
    }
}
